"""
A class and algorthims for solving Equalibrium problems on a network

Subclasses from a graph implementation that provides directional arcs and
weights, and provides a shortest path implementation.
"""
import graph

class EQGraph(graph.Graph):
    """
    This class extends a graph by providing functionality for equalibrium
    modelling

    Notably this class provides flows on the graph, and cost functions for the
    arcs
    """

    def __init__(self, nodesGraph=None, nodeData=None, posData = None,
                 flows = None, odData = None, costFunctions = None):
        """ Initialise a eqgraph with the additional information
        """
        graph.Graph.__init__(self,nodesGraph, nodeData)
        self.posData = {}
        self.flows = {}
        self.odData = {}
        self.costFunctions = {}
        if posData:
            self.posData = posData
        if odData:
            self.odData = odData
        if costFunctions:
            self.costFunctions = costFunctions

    def getArcCostFunction(self, n1, n2):
        """Return a cost function if one exists or return a zero function
        """
        if self.hasArc(n1,n2):
            if (n1,n2) in self.costFunctions:
                return self.costFunctions[(n1,n2)]
            else:
                #return a zero function
                return lambda f: graph.Graph.getArcw(self,n1,n2)
        else:
            raise KeyError, repr(n1) + " -> " + repr(n2)

    def getArcFlow(self, n1, n2):
        """Return a cost function if one exists or return a zero function
        """
        if self.hasArc(n1,n2):
            if (n1,n2) in self.flows:
                return self.flows[(n1,n2)]
            else:
                #return a zero
                return 0
        else:
            raise KeyError, repr(n1) + " -> " + repr(n2)

    def getOdData(self, n1, n2 = None):
        """Return the od data for a node or node pair
        """
        if n1 in self.o:
            if n2 is None:
                if n2 in self.odData:
                    return self.odData[n1]
                else:
                    return {}
            elif n2 in self.odData[n1]:
                return self.odData[n1][n2]
            else:
                return {}
        else:
            raise KeyError, repr(n1)

    def find_path_weight(self,path):
        """Determines the weight of a path on graph g by finding the sum of
        all the arc wieghts

        @param g : a graph
        @param path : a list of arcs
        @return weight: the sum of the wieghts

        example
        >>> g = EQGraph({0:{1:1}, 1:{2:1}, 2:{0:1} })
        >>> path = Path([0,1,2,0])
        >>> g.find_path_weight(path)
        3

        """
        weight = 0
        path = Path(path)
        for a in path.arcs():
            weight += self.getArcw(*a)
        return weight

    def getArcw(self,n1,n2):
        """ find the weights that correspond to the flows
        """
        f = self.getArcFlow(n1,n2)
        c = self.getArcCostFunction(n1,n2)
        return c(f)
        

    def updateArcWeights(self):
        """Changes the weights of the graph based on the flows and cost functions

        @param g             : a graph
        @param flows         : a dictionary flows for each arc
        @return costfunctions: a dictionary of cost functions 
        @return path_flows   : flows along the paths
        @return arc_flows    : flows along the arcs
        """
        for n1,n2 in self.xarcs():
            f = self.getArcFlow(n1,n2)
            c = self.getArcCostFunction(n1,n2)
            if c:
                self.setArcw(n1,n2,c(f))

    def find_equalibriation_function(self,path1,path2):
        """ determines the equalibriation function for two paths by looking at their
        current flow and increasing the flow along path1 by f and reducing the flow
        along path2 by f

        @param g: a graph with costfunctions and flows along the arcs
        @param path1: a list of arcs along a path flow will increase allong this
            path by f
        @param path2: a list of arcs along a path flow will decrease along this
            path by f
        @return fx: a function showing the diference between the two path lengths
        
        >>> g = EQGraph({0:{1:1}, 1:{2:1}, 2:{0:1} })
        >>> g.costFunctions[(0,1)] = lambda x:x**2
        >>> g.costFunctions[(1,2)] = lambda x:x+2
        >>> g.costFunctions[(2,0)] = lambda x:x*2
        >>> path1 = Path([0,1,2,0])
        >>> path2 = Path([])
        >>> fx = g.find_equalibriation_function(path1,path2)
        >>> fx(8.0) == 8.0**2 + 8.0 + 2 + 8.0*2
        True
        >>> fx = g.find_equalibriation_function(path1,path1)
        >>> fx(8.0) == 0.0
        True
        >>> g.flows = dict([(a,1) for a in g.arcs()])
        >>> fx = g.find_equalibriation_function(path1,path2)
        >>> fx(8.0) == (9.0**2 + 9.0 + 2 + 9.0*2)
        True
     
        """
        #note that common arcs can not be included    
        fpath1 = lambda f: sum([(self.getArcCostFunction(*a))(self.getArcFlow(*a)+f)
                                for a in path1.arcs() if a not in path2.arcs()])
        fpath2 =  lambda f: sum([(self.getArcCostFunction(*a))(self.getArcFlow(*a)-f)
                                 for a in path2.arcs() if a not in path1.arcs()])
        fx = lambda f : fpath1(f) - fpath2(f)
        return fx

    def all_or_nothing_assignment(self, odpairs = None, debug = False):
        """
        finds initial flows on a graph (all flows along the shortest path)
        
        @param g             : a graph
        @param odpairs       : a dictionary of dictionarys of the od pairs
        @return paths        : paths used
        @return path_flows   : flows along the paths
        @return arc_flows    : flows along the arcs
        
        >>> g = EQGraph({0:{1:1}, 1:{2:1}, 2:{0:1} })
        >>> od = {0:{1:1, 2:1}, 1:{0:1, 2:1}, 2:{0:1, 1:1}}

        >>> paths, path_flows, arc_flows = g.all_or_nothing_assignment(od)
        >>> path_flows[(0,1,2)]
        1
        >>> arc_flows[(0,1)]
        3
        """
        if odpairs is None:
            odpairs = self.odData
        #find shortest paths from all nodes in od
        arc_flows = {}
        path_flows = {}
        paths = []
        odpaths = {}
        #set the flows to zero
        for a in self.xarcs():
                arc_flows[a] = 0
        for orig in odpairs:
            dist, pred = self.shortestPaths(orig)
            #build paths and increment flows
            odpaths[orig] = {}
            for dest, flow in odpairs[orig].items():
                #construct the path
                path = [dest]
                node = dest
                while node in pred:
                    node = pred[node]
                    arc_flows[(node,path[-1])] = arc_flows[(node,path[-1])] + flow
                    path.append(node)
                path.reverse()
                path = Path(path)
                path_flows[path] = flow
                paths.append(path)
                odpaths[orig][dest] = [path]
        self.flows = arc_flows
        return odpaths, path_flows, arc_flows

    def iterativeArcFlowUpdate(self,newflow, w):
        """ a weighted update of the flows

        example
        >>> g = EQGraph({0:{1:1, 2:1}, 1:{2:1}, 2:{0:1} })
        >>> g.flows = {(0,1):1.0, (0,2):0.0}
        >>> new = {(0,1):0, (0,2):1.0}
        >>> g.iterativeArcFlowUpdate(new, 0.1)
        >>> print g.getArcFlow(0,1)
        0.9
        >>> print g.getArcFlow(0,2)
        0.1
        """
        for a,f in newflow.items():
            self.flows[a] = (1-w)*self.getArcFlow(*a) + w*f



    def find_equalibrium_flows(self, od = None, debug = False, EPS=10**-6,
                                iterprint = False):
        """finds equalibrium flows for the given graph and od information

        @param g: a graph with cost functions and flows defined on it
        @param od: origin destination information in the form of a dictionary
            of dictionaries

        example:
        >>> g = EQGraph({0:{1:1, 2:1, 3:1}, 1:{4:0}, 2:{4:0}, 3:{4:0} })
        >>> g.odData = {0:{4:10}}
        >>> g.costFunctions = {(0,1): BPR_function_generator(10.0,0.15,4.0,2.0),
        ...                   (0,2): BPR_function_generator(20.0,0.15,4.0,2.0),
        ...                   (0,3): BPR_function_generator(25.0,0.15,4.0,2.0)
        ...                   }
        >>> odpaths, path_flows, arc_flows = g.find_equalibrium_flows()
        >>> abs(arc_flows[(0,1)] - 4.15) <= 0.01
        True
        >>> abs(arc_flows[(0,2)] - 3.12) <= 0.01
        True
        >>> abs(arc_flows[(0,3)] - 2.72) <= 0.01
        True

        Braees Paradox from http://www.bgu.ac.il/~bargera/tntp/
        >>> g = EQGraph({1:{3:1, 4:1}, 3:{2:1, 4:1}, 4:{2:1}})
        >>> g.odData = {1:{2:6}}
        >>> g.costFunctions = g.costFunctions = {(1,3): BPR_function_generator(10**-8,10**9,1.0,1.0),
        ...                 (1,4): BPR_function_generator(50,0.02,1.0,1.0),
        ...                 (3,2): BPR_function_generator(50,0.02,1.0,1.0),
        ...                 (3,4): BPR_function_generator(10,0.1,1.0,1.0),
        ...                 (4,2): BPR_function_generator(10**-8,10**9,1.0,1.0) }
        >>> odpaths, path_flows, arc_flows = g.find_equalibrium_flows()
        
        """
        import scipy.optimize

        if od is None:
            od = self.odData
        #find an initial solution
        self.updateArcWeights()
        if debug: print "find Initial Assignment"
        odpaths, path_flows, arc_flows = self.all_or_nothing_assignment(od)
        self.flows = arc_flows
        terminated = False
        if not od: #empty od data
            return odpaths, path_flows, self.flows
        iter = 0
        while not terminated:
            if debug: print "iteration ", iter
            self.updateArcWeights()
            terminated = True
            for orig, destinations in od.items():
                for dest in destinations.keys():
                    self.updateArcWeights()
                    spath = Path(self.shortestPath(orig,dest))
                    #find the longest used path
                    lpath = max(odpaths[orig][dest],
                                 key = lambda p: self.find_path_weight(p))
                    if abs(self.find_path_weight(lpath) - self.find_path_weight(spath)) <= EPS:
                        #paths are of equal length
                        continue
                    eqf = self.find_equalibriation_function(spath,lpath)
                    #find the root
                    if debug: print "spath ", spath, self.find_path_weight(spath)
                    if debug: print "lpath ", lpath, self.find_path_weight(lpath)
                    
                    if debug: print "max flow move ", path_flows[lpath], eqf(path_flows[lpath]),eqf(0)
                    if eqf(path_flows[lpath]) <=EPS:
                        #move all of the flow
                        x = path_flows[lpath]
                    else:
                        x = scipy.optimize.brentq(eqf, a = 0, b = path_flows[lpath])
                    if abs(x) >= EPS:
                        #increment the arc flows
                        for a in spath.arcs():
                            self.flows[a] = self.flows[a] + x
                        for a in lpath.arcs():
                            self.flows[a] = self.flows[a] -x
                        # add the path
                        if spath not in odpaths[orig][dest]:
                            odpaths[orig][dest].append(spath)
                        #increment the flows
                        if spath not in path_flows:
                            path_flows[spath] = x
                        else:
                            path_flows[spath] = path_flows[spath] + x
                        path_flows[lpath] = path_flows[lpath] - x
                        if abs(path_flows[lpath])<=EPS:
                            #remove the path
                            odpaths[orig][dest].remove(lpath)
                        #if debug: print "Flows ", self.flows
                        terminated = False
            iter +=1
            if iterprint: print iter
                        
        return odpaths,path_flows, self.flows

    def find_total_cost(self):
        """ finds the total cost of the current flow solution
        """
        self.updateArcWeights()
        return sum([self.getArcCostFunction(*a)(self.getArcFlow(*a))*
                    self.getArcFlow(*a)
                    for a in self.xarcs()])

    def plot2d(self, **kwargs):
        """plots the graph in 2 dimensions using posData as locations
        """

        return graph.Graph.plot2d(self,posattr="posData",**kwargs)


class Path(tuple):
    """a class for storing paths as a list of nodes
    """
    def __init__(self,nodes):
        tuple(nodes)

    def arcs(self):
        """ iterate though the nodes to find the arcs
        """
        return [(self[i],self[i+1]) for i in range(len(self)-1)]

def find_path_from_pred(node,predessors):
    """finds a path from a predessor vector

    """
    path = [node]
    while node in predessors:
        node = predessors[node]
        path.append(node)
    path.reverse()
    return Path(path)

def BPR_function_generator(T0,alpha,beta,Qmax):
    """Generates the BPR function for link times in a traffic assignment

    given as $T_Q=T_0\[1+\alpha\(\frac{Q}{Q_{max}}\)^\beta\]
    @param T0: Free flow link time
    @param alpha:
    @param beta:
    @param Qmax: The capacity of the link
    """

    return lambda Q: BPR_function(Q,T0,alpha,beta,Qmax) 

def BPR_function(Q,T0,alpha,beta,Qmax):
    """The BPR function for link times in a traffic assignment

    given as $T_Q=T_0\[1+\alpha\(\frac{Q}{Q_{max}}\)^\beta\]
    @param Q: flow
    @param T0: Free flow link time
    @param alpha:
    @param beta:
    @param Qmax: The capacity of the link
    """
    if Q<=0.0:
        return T0
    else:
        return T0*(1+alpha*(Q/Qmax)**beta)

def _test():
    import doctest
    doctest.testmod()

def main():
     #Braees Paradox from http://www.bgu.ac.il/~bargera/tntp/
     g = EQGraph({1:{3:1, 4:1}, 3:{2:1}, 4:{2:1}})
     g.posData = {1:(0,0), 2:(2,0), 3:(1,-1), 4:(1,1)}
     g.odData = {1:{2:6}}
     g.costFunctions = {(1,3): BPR_function_generator(10**-8,10**9,1.0,1),
                         (1,4): BPR_function_generator(50,0.02,1.0,1),
                         (3,2): BPR_function_generator(50,0.02,1.0,1),
                         (3,4): BPR_function_generator(10,0.1,1.0,1),
                         (4,2): BPR_function_generator(10**-8,10**9,1.0,1) }
     #print g.plot2d()
     g.find_equalibrium_flows()
     print "total cost without arc ", g.find_total_cost()
     #g.plot2d(flows ="getArcw")
     #g.plot2d(flows ="getArcFlow")
     g.addArc(3,4,1)
     g.find_equalibrium_flows(debug=True)
     print "total cost with arc ", g.find_total_cost()
     #g.plot2d(flows ="getArcFlow")
     #g.plot2d(flows ="getArcw")
     print "Weights: ", g.arcsw()
     print "Flows: ", g.flows


if __name__ == "__main__":
    _test()
    main()        
